// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu, jhkim

#include "coin2/strategy/linear_model/agg.h"

#include <algorithm>

#include "coin2/strategy/linear_model/util.h"

namespace coin2::strategy::linear_model {

/* ------------------------------ Order Protos ------------------------------ */
using coin::proto::OrderDirection;
using coin::proto::OrderDuration;
using coin::proto::OrderType;
typedef SignalRMSEManager::stored_value_type signal_error_type;

void LmAggLogic::onFeed(
    const Symbol* symbol,
    const FeedUpdate& upd,
    const ::fastfeature::FeedMessage& feed_msg,
    LinearModel* linear_model,
    int64_t* lm_feed_id,
    const int64_t* global_feed_id,
    LinearModelDelegateProto* delegate,
    bool simulation) {
  int64_t timestamp_override = upd.timestamp();
  std::optional<double> midp_return;
  lholder_.UpdateFeed(symbol, upd, feed_msg);
  if (!lholder_.HasBookAndPos()) return;

  auto ask0p = lholder_.last_book->Ask0()->price;
  auto bid0p = lholder_.last_book->Bid0()->price;
  double midp = 0.5 * (ask0p + bid0p);

  if (use_signal_rmse_maker_ || use_signal_rmse_taker_) {
    signal_rmse_manager_.SetSimulation(
        simulation,
        lm_agg_config_.signal_rmse_disable_protolog_for_sim());

    signal_rmse_manager_.UpdateBook(timestamp_override, lholder_.last_book);
    signal_rmse_manager_.UpdateTimeOnly(timestamp_override);
    if (signal_rmse_manager_.GetSignalRMSE(
            timestamp_override,
            &lm_agg_config_,
            signal_rmses_)) {
      signal_rmse_ready_ = true;
    }
  }

  double signal_rmse_maker = signal_rmses_[signal_error_type::SIGNAL_ERROR_MAKER];

  TryTrigger(
      symbol,
      upd,
      feed_msg,
      linear_model,
      lm_feed_id,
      global_feed_id,
      delegate,
      timestamp_override,
      &midp_return);

  if (midp_return && lm_agg_config_.use_cancel_model()) {
    if (use_signal_rmse_maker_) {
      signal_rmse_manager_.UpdateSignal(timestamp_override, *midp_return);
      if (!signal_rmse_ready_) {
        TryClear(symbol, upd, feed_msg);
        return;
      }
    }

    *midp_return *= lm_agg_config_.signal_multiplier();
    const auto* ask_inside_order =
        lholder_.oe->GetInsideOrder(*lholder_.product_holder, OrderDirection::DIRECTION_SELL);
    const auto* bid_inside_order =
        lholder_.oe->GetInsideOrder(*lholder_.product_holder, OrderDirection::DIRECTION_BUY);

    double signal_lb;
    double signal_ub;
    if (!pred_interval_.GetNSigmaBound(
            *midp_return,
            lm_agg_config_.nsigma_edge(),
            &signal_lb,
            &signal_ub,
            lm_agg_config_.nsigma_edge_cap())) {
      return;
    }

    double current_stack = lholder_.oe->GetCurrentStack(*lholder_.product_holder);

    double skew_edge_bps = 0;
    clear_edge_manager_.UpdatePositionInStack(timestamp_override, current_stack);
    skew_edge_bps = clear_edge_manager_.clear_lean_bps
        - current_stack * lm_agg_config_.lean_per_stack_bps();

    double ask_skew_edge_bps = std::max(skew_edge_bps, -5.);
    double bid_skew_edge_bps = std::min(skew_edge_bps, 5.);
    if (skew_edge_bps > 0) {
      signal_ub += ask_skew_edge_bps;
      signal_lb += bid_skew_edge_bps * lm_agg_config_.lean_close_multiplier();
    } else {
      signal_ub += ask_skew_edge_bps * lm_agg_config_.lean_close_multiplier();
      signal_lb += bid_skew_edge_bps;
    }

    double var_edge = std::max(
        cancel_edge_bps_,
        lm_agg_config_.signal_rmse_maker_multiplier() * signal_rmse_maker);

    double ask_inside_bound = midp * (1 + 1e-4 * (signal_ub + var_edge));
    double ask_outside_bound = ask0p * (1 + 1e-4 * lm_agg_config_.cancel_outside_bps());
    double ask_post_price = ask_inside_bound * (1 + posting_edge_bps_add_ * 1e-4);

    if (new_ask_fill_ && midp > prev_ask_fill_price_ &&
        prev_ask_fill_ts_ + def_edge_ns_ > timestamp_override) {
      sell_def_edge_add_ += lm_agg_config_.defensive_edge_add_bps();
      new_ask_fill_ = false;
    } else {
      if (prev_sell_def_edge_restored_ts_ + def_edge_ns_ < timestamp_override) {
        sell_def_edge_add_ -= lm_agg_config_.defensive_edge_sub_bps();
        sell_def_edge_add_ = std::max(sell_def_edge_add_, 0.);
        prev_sell_def_edge_restored_ts_ = timestamp_override;
      }
    }

    ask_post_price *= (1 + 1e-4 * sell_def_edge_add_);
    ask_inside_bound *= (1 + 1e-4 * sell_def_edge_add_);
    ask_outside_bound *= (1 + 1e-4 * sell_def_edge_add_);

    std::optional<double> ask_sticky_outside;
    if (lm_agg_config_.use_sticky_outside()) {
      ask_sticky_outside = ask_post_price * (1 + lm_agg_config_.sticky_outside_bps() * 1e-4);
    }
    if (ask_post_price <= ask0p * (1 + lm_agg_config_.cling_to_bbo_bps() * 1e-4)) {
      ask_inside_bound = std::min(ask_inside_bound, ask0p);
      ask_post_price = std::min(ask_post_price, ask0p);
    }

    bool ask_post = true;
    if (lm_agg_config_.use_post_block() &&
        signal_ub > lm_agg_config_.post_block_threshold()) {
      ask_post = false;
    }

    if (ask_inside_order) {
      ask_post = false;
      double ask_working_price = ask_inside_order->order_price();
      double ask_too_outside = midp * (1 + 1e-4 * lm_agg_config_.ignore_too_outside_bps());

      if (ask_working_price > ask_too_outside && ask_post_price > ask_too_outside) {
        // both too outside: noisy action. do nothing
        ask_post = false;
      } else if (ask_working_price < ask_inside_bound) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_SELL, timestamp_override);
        ask_post = lm_agg_config_.layering();
      } else if (ask_working_price > ask_outside_bound) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_SELL, timestamp_override);
        ask_post = lm_agg_config_.layering();
      } else if (ask_sticky_outside && ask_working_price > *ask_sticky_outside) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_SELL, timestamp_override);
        ask_post = true;
      }
    }
    if (ask_post &&
        lm_agg_config_.use_passive_post() &&
        prev_ask_post_time_ + post_cooltime_ < timestamp_override) {
      prev_ask_post_time_ = timestamp_override;
      ask_post_price = std::max(
          ask_post_price,
          lholder_.product_holder->product_info().tick_price().GetNextPrice(bid0p));
      ask_post_price = std::max(
          ask_post_price,
          bid0p * (1 + lm_agg_config_.min_edge_bps() * 1e-4));

      if (ask_post_price <= ask_outside_bound) {
        double rounded_price =
            lholder_.product_holder->product_info().tick_price().RoundPassively(
                ask_post_price, OrderDirection::DIRECTION_SELL);
        double raw_qty = lholder_.oe->GetConfig(*lholder_.product_holder).qty().max_lot_size();
        raw_qty *= lm_agg_config_.maker_qty_ratio();

        lholder_.oe->SubmitOrder(
            *lholder_.product_holder,
            rounded_price,  // OE does rounding, so this can be replaced by
                            // ExecutorConfig-RoundingPolicy
            raw_qty,
            OrderDirection::DIRECTION_SELL,
            lm_agg_config_.order_duration(),
            OrderType::LIMIT_ORDER,
            ask0p,
            leverage_rate_,
            true,  // post_only
            "passive_post",
            timestamp_override,
            false,
            lm_agg_config_.use_block_price_band());
      }
    }

    double bid_inside_bound = midp * (1 + 1e-4 * (signal_lb - var_edge));
    double bid_outside_bound = bid0p * (1 - 1e-4 * lm_agg_config_.cancel_outside_bps());
    double bid_post_price = bid_inside_bound * (1 - posting_edge_bps_add_ * 1e-4);

    if (new_bid_fill_ && midp < prev_bid_fill_price_ &&
        prev_bid_fill_ts_ + def_edge_ns_ > timestamp_override) {
      buy_def_edge_add_ += lm_agg_config_.defensive_edge_add_bps();
      new_bid_fill_ = false;
    } else {
      if (prev_buy_def_edge_restored_ts_ + def_edge_ns_ < timestamp_override) {
        buy_def_edge_add_ -= lm_agg_config_.defensive_edge_sub_bps();
        buy_def_edge_add_ = std::max(buy_def_edge_add_, 0.);
        prev_buy_def_edge_restored_ts_ = timestamp_override;
      }
    }

    bid_post_price *= (1 - 1e-4 * buy_def_edge_add_);
    bid_inside_bound *= (1 - 1e-4 * buy_def_edge_add_);
    bid_outside_bound *= (1 - 1e-4 * buy_def_edge_add_);

    std::optional<double> bid_sticky_outside;
    if (lm_agg_config_.use_sticky_outside()) {
      bid_sticky_outside = bid_post_price * (1 - lm_agg_config_.sticky_outside_bps() * 1e-4);
    }
    if (bid_post_price >= bid0p * (1 - lm_agg_config_.cling_to_bbo_bps() * 1e-4)) {
      bid_inside_bound = std::max(bid_inside_bound, bid0p);
      bid_post_price = std::max(bid_inside_bound, bid0p);
    }

    bool bid_post = true;
    if (lm_agg_config_.use_post_block() &&
        signal_lb < -lm_agg_config_.post_block_threshold()) {
      bid_post = false;
    }

    if (bid_inside_order) {
      bid_post = false;
      double bid_working_price = bid_inside_order->order_price();
      double bid_too_outside = midp * (1 - 1e-4 * lm_agg_config_.ignore_too_outside_bps());

      if (bid_working_price < bid_too_outside && bid_post_price < bid_too_outside) {
        // both too outside: noisy action. do nothing
        bid_post = false;
      } else if (bid_working_price > bid_inside_bound) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_BUY, timestamp_override);
        bid_post = lm_agg_config_.layering();
      } else if (bid_working_price < bid_outside_bound) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_BUY, timestamp_override);
        bid_post = lm_agg_config_.layering();
      } else if (bid_sticky_outside && bid_working_price < *bid_sticky_outside) {
        lholder_.oe->CancelAllOrders(
            *lholder_.product_holder, OrderDirection::DIRECTION_BUY, timestamp_override);
        bid_post = true;
      }
    }
    if (bid_post &&
        lm_agg_config_.use_passive_post() &&
        prev_bid_post_time_ + post_cooltime_ < timestamp_override) {
      prev_bid_post_time_ = timestamp_override;
      bid_post_price = std::min(
          bid_post_price,
          lholder_.product_holder->product_info().tick_price().GetPrevPrice(ask0p));
      bid_post_price = std::min(
          bid_post_price,
          ask0p * (1 - lm_agg_config_.min_edge_bps() * 1e-4));

      if (bid_post_price >= bid_outside_bound) {
        double rounded_price =
            lholder_.product_holder->product_info().tick_price().RoundPassively(
                bid_post_price, OrderDirection::DIRECTION_BUY);
        double raw_qty = lholder_.oe->GetConfig(*lholder_.product_holder).qty().max_lot_size();
        raw_qty *= lm_agg_config_.maker_qty_ratio();

        lholder_.oe->SubmitOrder(
            *lholder_.product_holder,
            rounded_price,  // OE does rounding, so this can be replaced by
                            // ExecutorConfig-RoundingPolicy
            raw_qty,
            OrderDirection::DIRECTION_BUY,
            lm_agg_config_.order_duration(),
            OrderType::LIMIT_ORDER,
            bid0p,
            leverage_rate_,
            true,  // post_only
            "passive_post",
            timestamp_override,
            false,
            lm_agg_config_.use_block_price_band());
      }
    }
  }
  lholder_.oe->onPeriodicOrderCheck(*lholder_.product_holder, timestamp_override);
  if (!simulation) {
    lholder_.LogSignal(upd.timestamp(), delegate);
  }
}


void LmAggLogic::TryTrigger(
    const Symbol* symbol,
    const FeedUpdate& upd,
    const ::fastfeature::FeedMessage& feed_msg,
    LinearModel* linear_model,
    int64_t* lm_feed_id,
    const int64_t* global_feed_id,
    LinearModelDelegateProto* delegate,
    int64_t timestamp_override,
    std::optional<double>* midp_return) {
  CHECK(lholder_.target_symbol);
  delegate->ClearProto();
  lholder_.TriggerModelWithLatencyRecording(linear_model, lholder_.target_symbol);

  /* -------------------------------------------------------------------------- */
  /*                               lm_agg_logic.py                              */
  /* -------------------------------------------------------------------------- */

  auto ask0p = lholder_.last_book->Ask0()->price;
  auto bid0p = lholder_.last_book->Bid0()->price;

  /* ----------------------------- pick up signal ----------------------------- */
  double signal = 0;
  double sign = 0;
  double bbo_price = 0;
  double bbo_price_pass = 0;
  OrderDirection direction = OrderDirection::DIRECTION_UNINITIALIZED;
  std::string model_id = "";

  double buy_pie_in_the_sky = 0;
  double sell_pie_in_the_sky = 0;

  if (!delegate->response_.buy_edge_model().empty()) {
    buy_pie_in_the_sky = delegate->response_.buy_edge_model().Get(0).signal();
  }
  if (!delegate->response_.sell_edge_model().empty()) {
    sell_pie_in_the_sky = delegate->response_.sell_edge_model().Get(0).signal();
  }

  double pie_in_the_sky = 0;
  double leaned_threshold_bps = threshold_bps_;

  double calc_stack = lholder_.oe->GetCurrentStack(*lholder_.product_holder);

  if (!delegate->response_.buy_model().empty()) {
    double midp_ret = delegate->response_.buy_model().Get(0).signal();
    *midp_return = midp_ret;
    double spread_bps = (ask0p - bid0p) / ask0p * 1e4;
    model_id = delegate->response_.buy_model().Get(0).tag();
    if (midp_ret > 0) {
      sign = 1;
      direction = OrderDirection::DIRECTION_BUY;
      bbo_price = ask0p;
      bbo_price_pass = bid0p;
      pie_in_the_sky = buy_pie_in_the_sky;
      signal = midp_ret - penalize_spread_ratio_ * spread_bps;
      leaned_threshold_bps = threshold_bps_ +
                             (calc_stack > 0) * calc_stack * lm_agg_config_.lean_per_stack_bps();
    } else {
      sign = -1;
      direction = OrderDirection::DIRECTION_SELL;
      bbo_price = bid0p;
      bbo_price_pass = ask0p;
      pie_in_the_sky = sell_pie_in_the_sky;
      signal = -midp_ret - penalize_spread_ratio_ * spread_bps;
      leaned_threshold_bps =
          threshold_bps_ +
          (calc_stack < 0) * (-calc_stack) * lm_agg_config_.lean_per_stack_bps();
    }
  }

  std::optional<double> maxpos_multiplier;
  // exceeds 10 bp?
  if (sign > 0 && signal > 10) {
    maxpos_multiplier = lm_agg_config_.maxpos_fevermode_multiplier();
  }

  // signal window
  pred_interval_.UpdateSignal(timestamp_override, signal);
  double signal_lb;
  double signal_ub;
  if (!pred_interval_.GetNSigmaBound(
        signal,
        lm_agg_config_.nsigma_edge(),
        &signal_lb,
        &signal_ub,
        lm_agg_config_.nsigma_edge_cap())) {
    return;
  }

  if (direction == OrderDirection::DIRECTION_UNINITIALIZED) {
    return;
  }

  leaned_threshold_bps = std::min<double>(leaned_threshold_bps, threshold_bps_ * 2);

  double signal_rmse_taker_buy = signal_rmses_[signal_error_type::SIGNAL_ERROR_TAKER_BUY];
  double signal_rmse_taker_sell = signal_rmses_[signal_error_type::SIGNAL_ERROR_TAKER_SELL];
  if (sign > 0) {
    leaned_threshold_bps = std::max(
        leaned_threshold_bps,
        lm_agg_config_.signal_rmse_taker_multiplier() * signal_rmse_taker_buy);
  } else if (sign < 0) {
    leaned_threshold_bps = std::max(
        leaned_threshold_bps,
        lm_agg_config_.signal_rmse_taker_multiplier() * signal_rmse_taker_sell);
  }

  /* ----------------------------- threshold check ---------------------------- */
  double midp = 0.5 * (ask0p + bid0p);
  // auto signal_to_bp_divisor = is_model_output_bps_ ? 1 : midp / 1e4;
  // auto output_bps = signal / signal_to_bp_divisor;

  // limit payup only
  pie_in_the_sky = std::min<double>(pie_in_the_sky, lm_agg_config_.pie_in_the_sky_cap_bps());
  pie_in_the_sky = std::min<double>(pie_in_the_sky, lm_agg_config_.payup_cap_bps());

  // pie in the sky:
  // model penalty of profit due to predicted missing fills
  if (signal_lb - pie_in_the_sky < leaned_threshold_bps) {
    return;
  }

  /* ---------------------------------- lean ---------------------------------- */
  // if (lean_bps_ > 0) {
  //   double neutralized_potential_pos =
  //   double norm = neutralized_potential_pos / (max_pos - min_pos) * 2
  // }

  /* ------------------------------- agg pricing ------------------------------ */

  double payup_bps =
      signal - threshold_bps_ * lm_agg_config_.threshold_payup_margin_ratio();

  payup_bps = std::min(payup_bps, lm_agg_config_.payup_cap_bps());

  /* --------------------------------- submit --------------------------------- */

  double raw_qty = lholder_.oe->GetConfig(*lholder_.product_holder).qty().max_lot_size();
  bool post_only = false;

  double raw_price = (pie_in_the_sky < 0) ?
      bbo_price + sign * pie_in_the_sky * (midp * 1e-4)
      : bbo_price + sign * payup_bps * (midp * 1e-4);
  double rounded_price =
      lholder_.product_holder->product_info().tick_price().RoundPassively(raw_price, direction);

  if (lm_agg_config_.use_bbo_insteadof_taker()) {
    post_only = true;
    raw_price = bbo_price_pass;
    raw_qty /= 2;
  } else if (pie_in_the_sky < 0) {
    post_only = true;
    raw_qty /= 2;
  } else {
    if (!is_l1_book(*lholder_.last_book)) {
      raw_qty = std::min(
          raw_qty, get_available_qty(*lholder_.last_book, rounded_price, direction));
    }
  }

  if (use_signal_rmse_taker_ && !signal_rmse_ready_) return;

  if (raw_qty > 0) {
    std::string extra_tag = fmt::format(
        "{{\"logic\": \"lm-agg\", \"model_id\": \"{}\", \"signal\": {}}}",
        model_id,
        signal);
    bool fire_taker = true;
    if (direction == OrderDirection::DIRECTION_SELL &&
        prev_sell_taker_time_ + taker_cooltime_ < timestamp_override) {
      prev_sell_taker_time_ = timestamp_override;
    } else if (
        direction == OrderDirection::DIRECTION_BUY &&
        prev_buy_taker_time_ + taker_cooltime_ < timestamp_override) {
      prev_buy_taker_time_ = timestamp_override;
    } else {
      fire_taker = false;
      SPDLOG_INFO("Skipped taker due to the cooltime");
    }
    if (fire_taker) {
      lholder_.oe->SubmitOrder(
          *lholder_.product_holder,
          rounded_price,  // OE does rounding, so this can be replaced by
                          // ExecutorConfig-RoundingPolicy
          raw_qty,
          direction,
          lm_agg_config_.order_duration(),
          OrderType::LIMIT_ORDER,
          bbo_price,
          leverage_rate_,
          post_only,
          extra_tag,
          timestamp_override,
          false,  // give up logging due to the overhead
          lm_agg_config_.use_block_price_band(),
          maxpos_multiplier);
    }
    SPDLOG_INFO(extra_tag);
  }
}

}  // namespace coin2::strategy::linear_model
